from paddleocr import PaddleOCR
import cv2
import numpy as np
import json


def draw_rectangle(image, box, color=(0, 255, 0), thickness=2):
    """
    绘制轴对齐的矩形框
    :param image: 原始图片
    :param box: 四边形的四个顶点坐标 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
    :param color: 框的颜色
    :param thickness: 框的厚度
    """
    # 将四边形转换为最小外接矩形
    rect = cv2.boundingRect(np.array(box, dtype=np.int32))
    x, y, w, h = rect

    # 绘制矩形框
    cv2.rectangle(image, (x, y), (x + w, y + h), color, thickness)


def draw_char_rectangles(image, rect, text, counter, thickness=2):
    """
    绘制字符的矩形框，并交替使用红、黄、蓝三色，同时标注全局序号
    :param image: 原始图片
    :param rect: 文本行的矩形框 (x, y, w, h)
    :param text: 识别文本
    :param counter: 全局序号计数器
    :param thickness: 框的厚度
    :return: 更新后的全局序号计数器和矩形框信息
    """
    x, y, w, h = rect
    char_width = w / len(text)  # 假设字符宽度相等

    # 定义红、黄、蓝三色
    colors = [
        (0, 0, 255),  # 红色
        (0, 255, 255),  # 黄色
        (255, 0, 0)  # 蓝色
    ]

    # 用于存储矩形框信息的字典
    rect_info = {}

    for i, char in enumerate(text):
        # 计算当前字符的矩形框
        char_x = int(x + i * char_width)
        char_y = y
        char_w = int(char_width)
        char_h = h

        # 选择颜色（红、黄、蓝交替）
        color = colors[counter % 3]

        # 绘制字符的矩形框
        cv2.rectangle(image, (char_x, char_y), (char_x + char_w, char_y + char_h), color, thickness)

        # 在矩形框内部的左上角标注全局序号
        text_position = (char_x + 5, char_y + 20)  # 左上角偏移量
        cv2.putText(image, str(counter), text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, thickness)

        # 记录矩形框信息
        rect_info[counter] = {
            "top_left": (char_x, char_y),  # 左上角坐标
            "bottom_right": (char_x + char_w, char_y + char_h)  # 右下角坐标
        }

        # 更新全局序号计数器
        counter += 1

    return counter, rect_info


def main():
    # 初始化OCR模型
    ocr = PaddleOCR(
        use_angle_cls=True,  # 启用方向分类
        lang='ch',  # 中文模型
        det_db_thresh=0.1,  # 降低二值化阈值
        det_db_box_thresh=0.1,  # 降低框得分阈值
        det_db_unclip_ratio=1.5,  # 扩大框扩展比例
        det_db_box_size=3  # 设置最小框尺寸
    )

    # 读取图片
    image_path = './img/kongzi/a002.jpeg'
    try:
        image = cv2.imread(image_path)
        if image is None:
            raise FileNotFoundError(f"图片 {image_path} 未找到或无法读取")
    except Exception as e:
        print(f"图片读取失败: {e}")
        return

    # 执行OCR检测
    try:
        result = ocr.ocr(image_path, cls=True)
    except Exception as e:
        print(f"OCR检测失败: {e}")
        return

    # 初始化全局序号计数器
    counter = 0

    # 用于存储所有矩形框信息的字典
    all_rect_info = {}

    # 遍历所有检测到的文本行
    for line in result:
        for word_info in line:
            box = word_info[0]  # 提取坐标 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
            text = word_info[1][0]  # 提取识别文本

            # 绘制文本行的矩形框（绿色）
            draw_rectangle(image, box, color=(0, 255, 0), thickness=2)

            # 将四边形转换为最小外接矩形
            rect = cv2.boundingRect(np.array(box, dtype=np.int32))

            # 绘制字符的矩形框（红、黄、蓝交替）并标注全局序号
            counter, rect_info = draw_char_rectangles(image, rect, text, counter, thickness=2)

            # 将当前文本行的矩形框信息添加到全局字典中
            all_rect_info.update(rect_info)

    # 保存结果图片
    output_path = './img/kongzi/output_paddleocr_global_numbered_rectangles.jpg'
    cv2.imwrite(output_path, image)
    print(f"结果图片已保存到: {output_path}")

    # 保存矩形框信息为 JSON 文件
    json_path = './img/kongzi/rectangles_info.json'
    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(all_rect_info, f, ensure_ascii=False, indent=4)
    print(f"矩形框信息已保存到: {json_path}")


if __name__ == "__main__":
    main()
